In `@src/DistributedTraining/ActivationCheckpointConfig.cs`:
- Around line 52-75: The CheckpointEveryNLayers and MaxActivationsInMemory
properties lack validation and can lead to division-by-zero or invalid negative
values; add input checks in ActivationCheckpointConfig: enforce
CheckpointEveryNLayers > 0 (unless a special mode explicitly allows 0—if so
document and guard usage) and enforce MaxActivationsInMemory >= 0. Implement
this by adding validation logic (either in property setters for
CheckpointEveryNLayers and MaxActivationsInMemory, or a public Validate() method
on ActivationCheckpointConfig that throws
ArgumentException/ArgumentOutOfRangeException with clear messages) and ensure
callers construct/validate instances (e.g., call Validate() from the constructor
or factory) so invalid configs fail fast.

In `@src/DistributedTraining/GPipeSchedule.cs`:
- Around line 41-57: In GetSchedule, validate numStages (and numMicroBatches)
before checking stageId so you don't throw ArgumentOutOfRangeException for
stageId when numStages is invalid; move the numStages <= 0 check (and
numMicroBatches <= 0 if desired) above the stageId bounds check in the
GetSchedule method and keep the stageId validation afterwards, throwing
ArgumentException for numStages and ArgumentOutOfRangeException for an invalid
stageId.

In `@src/DistributedTraining/LoadBalancedPartitionStrategy.cs`:
- Around line 55-141: The constructor and BuildLayerSizes currently conflate a
single-element _layerBoundaries array with "auto-detect" behavior and do not
validate ordering/ranges; to fix, change the int[] ctor
(LoadBalancedPartitionStrategy(int[] layerBoundaries,...)) to validate that
layerBoundaries is non-null, has length >=1, contains strictly increasing,
non-negative values and that the last boundary < totalParameters (or at least
document/validate later), and throw ArgumentException for invalid input; keep
auto-detect behavior only for the other ctor LoadBalancedPartitionStrategy(int
estimatedLayerSize,...) by adding a private flag like _isAutoDetect (set in the
estimatedLayerSize ctor) and have BuildLayerSizes check _isAutoDetect (not
_layerBoundaries.Length==1) to generate synthetic layers, and in BuildLayerSizes
also validate boundaries are sorted and within range and compute sizes using
consecutive boundary differences so parameters aren’t silently dropped.

In `@src/DistributedTraining/PipelineParallelModel.cs`:
- Around line 216-224: The checkpointing implementation is incomplete and leaves
dead state; either remove the partial logic or make it fail-fast. Update the
code paths that reference RecomputeStrategy and CheckpointFirstLayer so that if
any checkpointing mode is selected the PipelineParallelModel throws a clear
NotImplementedException (or ArgumentException) at construction or before the
forward loop; remove or stop populating unused microBatchOutputs and only keep
microBatchInputs and _checkpointedActivations if they are actively used,
otherwise delete those fields and related writes to eliminate dead state; ensure
any remaining checkpoint-related flags are documented and guarded so no silent
partial behavior runs in production.
- Around line 226-243: The loop is reusing the same input/output for every
micro‑batch which corrupts gradients when _microBatchSize > 1; update the code
that builds microBatchInputs/microBatchOutputs and checkpointing to index into
the per‑microbatch collections using op.MicroBatchIndex (i.e., obtain stageInput
= <microbatch-list>[op.MicroBatchIndex] or call a proper slice helper instead of
reusing the top‑level input), store checkpointed activation into
_checkpointedActivations[op.MicroBatchIndex], and similarly ensure the
expectedOutput/loss lookup uses expectedOutputList[op.MicroBatchIndex] (or fail
fast if only scalar inputs are supported). Locate and fix references around
GetStageInput, ShouldCheckpointActivation, _checkpointedActivations,
microBatchInputs, WrappedModel.Predict, microBatchOutputs and the expectedOutput
usage so every microbatch uses its own indexed input/output.
- Around line 87-139: Public constructor and properties (PipelineParallelModel,
Schedule, PartitionStrategy, CheckpointConfig, CheckpointConfig) expose knobs
that bypass the intended facade; either make PipelineParallelModel internal and
keep the public surface via AiModelBuilder/AiModelResult, or keep the class
public but make those properties/constructor overloads internal/private and add
corresponding configuration entry points on AiModelBuilder that set
schedule/partition/checkpoint before building. Update visibility for the
constructor and/or Schedule/PartitionStrategy/CheckpointConfig properties or
move construction logic behind AiModelBuilder methods (e.g.,
AddPipelineSchedule, WithPartitionStrategy, WithCheckpointConfig) so external
users only interact through AiModelBuilder/AiModelResult.
- Around line 246-283: The tag math currently uses overlapping namespaces
(SendActivationsForward using tag: op.MicroBatchIndex * 10 and gradient
Send/Receive using tag: 1000 + op.MicroBatchIndex), which can collide for large
microBatchIndex; introduce dedicated constants (e.g., ACTIVATION_TAG_BASE and
GRADIENT_TAG_BASE or ACTIVATION_TAG_MULTIPLIER and GRADIENT_TAG_BASE) and
replace occurrences in SendActivationsForward, the gradient Send/Receive calls
(where tag is 1000 + op.MicroBatchIndex) and the other region mentioned (lines
~340-370) so activations use ACTIVATION_TAG_BASE + op.MicroBatchIndex and
gradients use GRADIENT_TAG_BASE + op.MicroBatchIndex (or multiply
microBatchIndex by a large non-overlapping multiplier) to guarantee
non‑overlapping tag ranges.
- Around line 171-177: The code uses
_partitionStrategy.ComputePartition(totalParams, _numStages) and immediately
indexes partitions[_stageId]; validate the returned partitions before indexing
by checking partitions is not null, partitions.Length == _numStages, and that
the entry for partitions[_stageId] has non‑negative StartIndex and Size and that
StartIndex + Size <= totalParams; if any check fails, throw an informative
exception (or fall back to a safe default partitioning) rather than proceeding
to assign ShardStartIndex and ShardSize from an invalid partition.

In `@src/Interfaces/IAiModelBuilder.cs`:
- Around line 769-781: The change to the public interface
IAiModelBuilder<T,TInput,TOutput>.ConfigureDistributedTraining alters its
signature and will break external implementers; restore the original interface
method signature (keep ConfigureDistributedTraining as it was) and move the new
pipeline-specific parameters into a non-breaking alternative such as: add an
overload on the concrete AiModelBuilder class or introduce a
PipelineDistributedOptions/DistributedTrainingOptions object that the facade
AiModelBuilder exposes (or add a new method ConfigurePipelineDistributedTraining
on AiModelBuilder) so external implementations of IAiModelBuilder are unaffected
while still supporting pipelineSchedule, pipelinePartitionStrategy,
pipelineCheckpointConfig and pipelineMicroBatchSize.

In `@src/DistributedTraining/Interleaved1F1BSchedule.cs`:
- Around line 95-97: In Interleaved1F1BSchedule.cs, remove the redundant "/ 1"
from the computation of numWarmupForwards (currently: int numWarmupForwards =
Math.Min((totalVirtualStages - 1 - stageId) / 1, numMicroBatches *
_virtualStagesPerRank);); change the left-side expression to simply
(totalVirtualStages - 1 - stageId) so Math.Min compares that value directly with
numMicroBatches * _virtualStagesPerRank, preserving behavior but eliminating the
no-op division.
- Around line 70-84: The validation currently checks stageId before ensuring
numStages is positive; update the checks in Interleaved1F1BSchedule (the
constructor or validation block that references stageId, numStages, and
numMicroBatches) to validate numStages > 0 first, then validate that stageId is
within [0, numStages-1], and keep the numMicroBatches > 0 check as is; reorder
the checks so ArgumentException for numStages comes before the
ArgumentOutOfRangeException for stageId.
- Around line 103-104: Remove the dead tracking arrays forwardCount and
backwardCount from the Interleaved1F1BSchedule class: delete their declarations
(var forwardCount/backwardCount) and remove all statements that mutate them
(increments/assignments where forwardCount[...]++ or backwardCount[...]++
appear), since their values are never read; ensure no other code references
these identifiers and run tests/compile to confirm no remaining references.

In `@src/DistributedTraining/LoopedBFSSchedule.cs`:
- Around line 73-87: The validation for stageId is performed before ensuring
numStages is valid in LoopedBFSSchedule (constructor or validation block); move
the check "if (numStages <= 0) { throw new ArgumentException(...,
nameof(numStages)); }" to run before the "if (stageId < 0 || stageId >=
numStages) { throw new ArgumentOutOfRangeException(nameof(stageId), ...); }"
check so stageId comparisons are only done when numStages is positive, and keep
the existing validation for numMicroBatches as-is (numMicroBatches <= 0).
- Around line 105-106: In LoopedBFSSchedule.cs inside the method computing
vStage, remove the two dead local variables isFirstLoop and isLastLoop (they are
declared as bool isFirstLoop = vStage == 0; and bool isLastLoop = vStage ==
_virtualStagesPerRank - 1;) since they are never used; alternatively, if special
warmup/cooldown logic was intended, replace their declarations with the actual
handling logic referencing vStage and _virtualStagesPerRank, but do not leave
unused locals behind.

In `@src/DistributedTraining/OneForwardOneBackwardSchedule.cs`:
- Around line 52-66: The validation currently checks stageId bounds before
ensuring numStages is positive; in OneForwardOneBackwardSchedule
(constructor/initializer) move the numStages <= 0 check to run before the
stageId range check so you validate the container size first, then verify
stageId is within 0..numStages-1; keep the numMicroBatches <= 0 check as-is and
preserve the same exception types/messages.

In `@src/DistributedTraining/PipelineParallelModel.cs`:
- Around line 209-210: The schedule returned by _schedule.GetSchedule(_stageId,
_numStages, _microBatchSize) may include invalid micro-batch indices; before
executing scheduleOps, validate every op in scheduleOps (use the existing
scheduleOps variable and the types it contains) to ensure any micro-batch index
field is within [0, _microBatchSize - 1] (and optionally that any target
stage/index fields are within valid stage range 0.._numStages-1); if any entry
is out of bounds, throw an ArgumentException or similar with a clear message
identifying the offending op and the expected bounds so invalid
externally-injected IPipelineSchedule implementations fail fast.

In `@src/DistributedTraining/ZeroBubbleH1Schedule.cs`:
- Around line 112-122: Remove the redundant "- 0" from the conditional in
ZeroBubbleH1Schedule: the check currently uses "backwardInputIdx - 0", which is
a no-op; update the if condition in the block that adds a new PipelineOperation
(the variables backwardWeightIdx, backwardInputIdx, and numMicroBatches, and the
creation of a PipelineOperation with Type =
PipelineOperationType.BackwardWeight) to compare backwardWeightIdx directly
against backwardInputIdx (i.e., use "backwardWeightIdx < backwardInputIdx &&
backwardWeightIdx < numMicroBatches") so the logic remains unchanged but the
code is cleaned up.
- Around line 39-53: The validation currently checks stageId bounds before
verifying numStages is positive; move the numStages check (the throw for
numStages <= 0) to run before the stageId range check so that stageId is not
compared against an invalid numStages, and keep the existing numMicroBatches > 0
check as-is; update the validation order in the ZeroBubbleH1Schedule
constructor/method (references: stageId, numStages, numMicroBatches) and apply
the same reordering to the other schedule implementations that use the same
checks.

In `@src/DistributedTraining/ZeroBubbleH2Schedule.cs`:
- Around line 37-51: The validation currently checks stageId before numStages
which can produce misleading errors; in the ZeroBubbleH2Schedule constructor (or
the method that validates inputs), move the check for numStages (numStages <= 0)
to run before the stageId range check, then keep the stageId check (stageId < 0
|| stageId >= numStages) and the numMicroBatches check (numMicroBatches <= 0)
as-is so errors are accurate and deterministic.

In `@src/DistributedTraining/ZeroBubbleVSchedule.cs`:
- Around line 1-263: Extract the repeated parameter checks in
ZeroBubbleVSchedule.GetSchedule into a shared validation helper and call it from
GetSchedule: move the three checks (numStages > 0, stageId in [0,numStages-1],
numMicroBatches > 0) into a new internal static
ScheduleValidation.ValidateGetScheduleParameters(int stageId, int numStages, int
numMicroBatches) and replace the inline checks at the top of
ZeroBubbleVSchedule.GetSchedule with a single call to that helper; apply the
same change to the other schedule classes so all seven schedules use the common
ScheduleValidation helper to remove duplication and ensure consistency.
- Around line 51-65: The code in ZeroBubbleVSchedule.cs validates stageId before
checking numStages, which can throw an ArgumentOutOfRangeException when
numStages is invalid; reverse the checks to validate numStages and
numMicroBatches first, then validate stageId. Extract a shared helper method
(e.g., ValidateScheduleParameters or ValidateStageParams) that takes numStages,
numMicroBatches and stageId and performs: (1) numStages > 0, (2) numMicroBatches
> 0, (3) 0 <= stageId < numStages, then call that helper from
ZeroBubbleVSchedule (and other schedule implementations) to ensure consistent
validation across the codebase.

In `@src/DistributedTraining/PipelineParallelModel.cs`:
- Around line 583-628: GetStageInput currently falls back to the original
micro-batch for non-first virtual stages on the same rank instead of using the
previous virtual stage's forward output; replace that fallback logic in
GetStageInput so that when virtualStageIndex > 0 and not receiving from a
previous rank you lookup the previous virtual stage's output from forwardOutputs
using the prior op key (opKey - 1 or the equivalent key construction used
elsewhere), convert that Vector<T> to TInput via
ConversionsHelper.ConvertVectorToInputWithoutReference<T, TInput>(...), and
return it; only if forwardOutputs does not contain the prior-stage output then
fall back to microBatches and otherwise throw the existing
InvalidOperationException.
- Around line 678-695: The ShouldCheckpointActivation method currently does an
opKey % _checkpointConfig.CheckpointEveryNLayers which can divide by zero; add a
guard that checks _checkpointConfig.CheckpointEveryNLayers > 0 before using the
modulo (e.g., if <= 0, treat checkpointing interval as disabled and return false
or surface a configuration error), updating ShouldCheckpointActivation to first
validate _checkpointConfig.CheckpointEveryNLayers and avoid the modulo when it
is non-positive.
- Around line 723-754: The nearest-checkpoint search in the recompute block can
pick checkpoints from other micro-batches because it only checks opKey against
_checkpointedActivations; change the search to restrict candidates to the
current micro-batch (e.g. start searchKey at microBatchIndex * V and only
consider keys in range microBatchIndex * V .. opKey-1) or refactor checkpoint
storage/lookup to use a composite key (microBatchIndex, virtualStageIndex) so
you only fetch a checkpoint from the same micro-batch; update the loop that sets
nearestCheckpointKey and the subsequent access of _checkpointedActivations to
use the new range or composite key check before converting checkpointVector and
running WrappedModel/Predict via ConversionsHelper.
- Around line 452-569: In SliceInputIntoMicroBatches and
SliceTargetIntoMicroBatches, stop silently duplicating data when conversion
fails or microBatchElements <= 0; instead throw a clear exception (e.g.,
ArgumentException or InvalidOperationException) indicating the provided
TInput/TOutput is not sliceable for the configured _microBatchSize (include
_microBatchSize and a brief context in the message). Replace the conversion
catch blocks and the microBatchElements <= 0 branches so they throw rather than
populate all slices, and ensure the exception type and message make it obvious
which method (SliceInputIntoMicroBatches or SliceTargetIntoMicroBatches) and
which parameter (input/target) caused the failure.

using AiDotNet.Interfaces;

In `@src/AiDotNet.csproj`:
- Around line 142-146: The project currently enables EmitCompilerGeneratedFiles
and sets CompilerGeneratedFilesOutputPath unconditionally; move these entries
into a configuration-gated PropertyGroup (e.g. PropertyGroup with
Condition="'$(Configuration)' == 'Debug'") or wrap them behind a dedicated
MSBuild property so they only apply in Debug (or when the property is explicitly
enabled), referencing the existing EmitCompilerGeneratedFiles and
CompilerGeneratedFilesOutputPath property names and the PropertyGroup element to
locate where to change.

In `@src/AiModelBuilder.cs`:
- Around line 3738-3796: The public ConfigurePipelineParallelism method
currently accepts non-positive microBatchSize which breaks downstream pipeline
scheduling; update ConfigurePipelineParallelism to validate microBatchSize > 0
and fail fast by throwing an ArgumentOutOfRangeException (or ArgumentException)
with a clear message when invalid, before assigning to _pipelineMicroBatchSize;
reference the method name ConfigurePipelineParallelism, parameter
microBatchSize, and field _pipelineMicroBatchSize when making the change.

In `@src/DistributedTraining/ActivationCheckpointConfig.cs`:
- Around line 56-105: The ArgumentOutOfRangeException calls inside the
CheckpointEveryNLayers setter and the MaxActivationsInMemory setter use
nameof(value) which yields "value"; change those to use the property names
instead (replace nameof(value) with nameof(CheckpointEveryNLayers) in the
CheckpointEveryNLayers setter and with nameof(MaxActivationsInMemory) in the
MaxActivationsInMemory setter) so the thrown exceptions reference the actual
property names and match .NET patterns.

In `@src/DistributedTraining/PipelineParallelModel.cs`:
- Around line 811-844: The recomputation block guarded by
_checkpointConfig.Enabled and _checkpointConfig.RecomputeStrategy !=
RecomputeStrategy.None (which uses _checkpointedActivations, ConversionsHelper,
WrappedModel.Predict, etc.) is dead because the PipelineParallelModel
constructor throws NotImplementedException for any RecomputeStrategy other than
None; either remove this entire block to avoid maintaining unreachable code or
retain it but add a clear comment/TODO above it explaining it is
placeholder/infrastructure for future recompute support and currently
unreachable due to the constructor guard so reviewers won’t treat it as active
logic.
- Around line 892-911: AccumulateGradients currently assumes accumulated.Length
== newGradients.Length; add a defensive check at the start of the method (after
the accumulated null check) that validates newGradients is non-null and that
accumulated.Length == newGradients.Length, and if not throw an ArgumentException
(or ArgumentNullException if newGradients is null) with a clear message
containing the two lengths and parameter names; keep the existing clone and
element-wise addition logic in AccumulateGradients unchanged.
- Around line 125-128: EstimatedBubbleFraction currently calls
_schedule.EstimateBubbleFraction(_numStages, _microBatchSize) but _numStages is
only set in OnBeforeInitializeSharding/InitializeSharding, so access before
initialization yields invalid results; update the EstimatedBubbleFraction getter
to guard against uninitialized state by checking _numStages (or an
initialization flag set in OnBeforeInitializeSharding/InitializeSharding) and
either throw a clear InvalidOperationException or return a safe default (e.g.,
0) when not initialized, referencing EstimatedBubbleFraction, _numStages,
OnBeforeInitializeSharding/InitializeSharding and
_schedule.EstimateBubbleFraction so reviewers can locate and apply the fix.
- Around line 206-234: The multi-stage branch that uses _partitionStrategy and
vsPartitions currently only checks count; add the same per-partition bounds
validation used in the single-stage path: for each entry returned by
ComputePartition(totalParams, _totalVirtualStages) verify StartIndex >= 0, Size
>= 0, and StartIndex + Size <= totalParams, and throw an
InvalidOperationException (including which partition index and offending values)
if any check fails so you catch out-of-range or negative partitions before they
are used by Array.Copy or later logic.

In `@src/DistributedTraining/ZeroBubbleH2Schedule.cs`:
- Around line 34-79: GetSchedule validates stageId but never uses it, producing
identical schedules for all stages; change the logic to make the schedule
stage-local by using stageId to decide which global timesteps this stage
actually executes: compute global forward timesteps and only add a
PipelineOperation.Forward when the global forward counter maps to this stage
(e.g., globalForwardIndex % numStages == stageId), set MicroBatchIndex using
that globalForwardIndex-to-microbatch mapping, and similarly for backward passes
use a globalBackwardIndex where you add a PipelineOperation.Backward only when
(globalBackwardIndex % numStages == stageId) (or the equivalent reverse-stage
mapping used by your other schedulers); adjust
numWarmupForwards/forwardIdx/steadyStateCount/backwardInputIdx/backwardWeightIdx
calculations so they operate on global counters and derive per-stage counts by
modular filtering, and ensure IsWarmup/IsCooldown flags are computed from the
global timeline so later stages don’t run work before inputs are available
(update GetSchedule, numWarmupForwards, forwardIdx, steadyStateCount,
backwardInputIdx, backwardWeightIdx and the PipelineOperation creation sites).

In `@src/DistributedTraining/ZeroBubbleVSchedule.cs`:
- Around line 185-209: The BackwardWeight ops are being marked IsCooldown = true
unconditionally; compute the isCooldown flag once per loop iteration like you do
for BackwardInput and reuse it for both BackwardWeight branches. Update the code
that builds PipelineOperation for BackwardWeight (the blocks that set Type =
PipelineOperationType.BackwardWeight and touch
MicroBatchIndex/backwardWeightCount0/backwardWeightCount1/VirtualStageIndex) to
set IsCooldown = isCooldown (the same local boolean computed earlier in the
loop) instead of hardcoding true, and ensure the increment of
backwardWeightCount0/backwardWeightCount1 remains unchanged.

In `@src/Interfaces/IAiModelBuilder.cs`:
- Line 789: Rename the ambiguous parameter microBatchSize to microBatchCount in
the IAiModelBuilder interface and any related method signatures (e.g.,
IAiModelBuilder.BuildModel / methods referencing microBatchSize) and update the
XML doc comment from "Number of micro-batches" to explicitly "Number of
micro-batches (micro-batch count for pipeline execution)". Ensure all usages,
XML <param> tags, and any implementing classes or callers are updated to use
microBatchCount to keep names consistent.